# Copyright (c) 2019-2021, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import print_function

import argparse
import logging
import os
import sys

import flask
from flask import Flask, render_template
from tornado.wsgi import WSGIContainer
from tornado.web import Application, FallbackHandler
from tornado.ioloop import IOLoop

from kaolin.experimental.dash3d.util import StreamingGeometryHelper, GeometryWebSocketHandler

logger = logging.getLogger(__name__)


def get_max_viewports(urlargs):
    default_val = 3
    try:
        res = int(urlargs.get("maxviews", default_val))
        res = max(1, min(8, res))
    except Exception as e:
        res = default_val
    return res


def create_server(logdir):
    """ Create the server, including websocket handler through tornado, and flask http server.

    Args:
        logdir (str): directory where Timelapse data is written.
    """
    # Helper for streaming geometry from the logdir
    helper = StreamingGeometryHelper(logdir)

    # Flask for HTTP
    _base_dir = os.path.dirname(__file__)
    _template_dir = os.path.join(_base_dir, 'templates')
    _static_dir = os.path.join(_base_dir, 'static')
    app = Flask('kaolin_dash3d',
                template_folder=_template_dir,
                static_url_path='/static',
                static_folder=_static_dir)

    @app.route('/')
    def index():
        helper.parser.check_for_updates()

        urlargs = dict(flask.request.args)
        max_viewports = get_max_viewports(urlargs)
        return render_template('home.html', logdir=helper.logdir,
                               nmeshes=min(helper.parser.num_mesh_items(), max_viewports),
                               npointclouds=min(helper.parser.num_pointcloud_items(), max_viewports),
                               urlargs=urlargs)

    # Tornado server to handle websockets
    container = WSGIContainer(app)
    server = Application([
        (r'/websocket/', GeometryWebSocketHandler, dict(helper=helper)),
        (r'.*', FallbackHandler, dict(fallback=container))
    ])
    return server


def run_main():
    aparser = argparse.ArgumentParser(
        description='NVIDIA Kaolin Tensorboard 3d visualizer for USD files generated during training.')
    aparser.add_argument('--logdir', action='store', type=str, required=True,
                         help='The vis folder generated by the Timelapse module.')
    aparser.add_argument('--log_level', action='store', type=int, default=logging.INFO,
                         help='Logging level, DEBUG: 10, INFO: 20, WARN: 30, ERROR: 40.')
    aparser.add_argument('--port', action='store', default=8080)
    args = aparser.parse_args()

    logging.basicConfig(level=args.log_level,
                        format='%(asctime)s|%(levelname)8s|%(name)15s| %(message)s',
                        handlers=[logging.StreamHandler(sys.stdout)])

    print(f'Dash3D server starting. Go to: http://localhost:{args.port}')
    server = create_server(args.logdir)
    server.listen(args.port)
    IOLoop.instance().start()


if __name__ == "__main__":
    run_main()
